import MCPClient from "./MCPClient";
import ChatOpenAI from "./ChatOpenAI";
import { logTitle } from "./utils";

export default class Agent {
  private mcpClients: MCPClient[];
  private llm: ChatOpenAI | null = null;
  private model: string;
  private systemPrompt: string;
  private context: string;
  private originalPrompt: string = "";
  private failureCount: Map<string, number> = new Map();
  private maxFailures: number = 3;
  private promptCorrectionEnabled: boolean = true;

  constructor(
    model: string,
    mcpClients: MCPClient[],
    systemPrompt: string = "",
    context: string = ""
  ) {
    this.mcpClients = mcpClients;
    this.model = model;
    this.systemPrompt = systemPrompt;
    this.context = context;
    this.originalPrompt = "";
    this.failureCount = new Map();
    this.maxFailures = 3;
    this.promptCorrectionEnabled = true;
  }

  async init() {
    logTitle("TOOLS");
    for await (const client of this.mcpClients) {
      await client.init();
    }
    const tools = this.mcpClients.flatMap((client) => client.getTools());
    this.llm = new ChatOpenAI(
      this.model,
      this.systemPrompt,
      tools,
      this.context
    );
  }

  async close() {
    for await (const client of this.mcpClients) {
      await client.close();
    }
  }

  async invoke(prompt: string) {
    if (!this.llm) throw new Error("Agent not initialized");

    // 存储原始提示词（仅在第一次调用时）
    if (!this.originalPrompt) {
      this.originalPrompt = prompt;
    }

    let response = await this.llm.chat(prompt);
    while (true) {
      if (response.toolCalls.length > 0) {
        for (const toolCall of response.toolCalls) {
          const mcp = this.mcpClients.find((client) =>
            client
              .getTools()
              .some((t: any) => t.name === toolCall.function.name)
          );
          if (mcp) {
            logTitle(`TOOL USE`);
            process.stdout.write(`Calling tool: ${toolCall.function.name}`);
            process.stdout.write(`Arguments: ${toolCall.function.arguments}`);

            try {
              const parsedArgs = JSON.parse(toolCall.function.arguments);
              const result = await mcp.callTool(
                toolCall.function.name,
                parsedArgs
              );
              process.stdout.write(`Result: ${JSON.stringify(result)}`);
              process.stdout.write(JSON.stringify(result));
              this.llm.appendToolResult(toolCall.id, JSON.stringify(result));
            } catch (error) {
              // 处理JSON解析错误和工具调用错误
              const errorMessage = (error as Error).message ?? String(error);
              process.stdout.write(
                `Tool ${toolCall.function.name} error: ${errorMessage}`
              );
              this.llm.appendToolResult(toolCall.id, `Error: ${errorMessage}`);

              // 记录工具调用失败
              await this.handleToolFailure(
                toolCall.function.name,
                errorMessage
              );
            }
          } else {
            this.llm.appendToolResult(toolCall.id, "Tool not found");
          }
        }
        // 工具调用后,继续对话
        response = await this.llm.chat();
        continue;
      }
      // 没有工具调用时，提示继续执行任务
      if (response.content && response.content.trim().length > 0) {
        // 如果有内容输出，继续对话以促进工具调用
        response = await this.llm.chat("请继续执行任务，使用相应的工具完成下一步操作。");
        continue;
      } else {
        // 没有内容输出，结束对话
        await this.close();
        return response.content;
      }
    }
  }

  /**
   * 处理工具调用失败，记录失败次数并在必要时触发提示词纠正
   */
  private async handleToolFailure(toolName: string, errorMessage: string) {
    if (!this.promptCorrectionEnabled) return;

    const currentCount = this.failureCount.get(toolName) || 0;
    this.failureCount.set(toolName, currentCount + 1);

    logTitle(`TOOL FAILURE TRACKING`);
    process.stdout.write(`Tool ${toolName} failed ${currentCount + 1} times`);

    if (currentCount + 1 >= this.maxFailures) {
      process.stdout.write(
        `Tool ${toolName} has failed ${this.maxFailures} times. Triggering prompt correction...`
      );
      await this.correctPrompt(toolName, errorMessage);
      // 重置失败计数
      this.failureCount.set(toolName, 0);
    }
  }

  /**
   * 智能提示词纠正方法
   */
  private async correctPrompt(failedToolName: string, lastError: string) {
    if (!this.llm) return;

    logTitle(`PROMPT CORRECTION`);
    process.stdout.write(`Analyzing failures for tool: ${failedToolName}`);

    // 获取工具的详细信息
    const toolInfo = this.getToolInfo(failedToolName);

    // 构建纠正提示词
    const correctionPrompt = this.buildCorrectionPrompt(
      failedToolName,
      lastError,
      toolInfo
    );

    // 使用LLM分析并生成改进的提示词
    const correctedInstructions = await this.generateCorrectedInstructions(
      correctionPrompt
    );

    // 应用纠正后的指令
    await this.applyCorrectedInstructions(correctedInstructions);
  }

  /**
   * 获取工具信息
   */
  private getToolInfo(toolName: string) {
    for (const client of this.mcpClients) {
      const tool = client.getTools().find((t: any) => t.name === toolName);
      if (tool) {
        return {
          name: tool.name,
          description: tool.description,
          inputSchema: tool.inputSchema,
        };
      }
    }
    return null;
  }

  /**
   * 构建纠正提示词
   */
  private buildCorrectionPrompt(
    toolName: string,
    error: string,
    toolInfo: any
  ): string {
    return `
作为一个智能提示词优化专家，请分析以下情况并提供改进建议：

**原始任务：**
${this.originalPrompt}

**失败的工具：** ${toolName}
**最近的错误：** ${error}

**工具详细信息：**
- 名称：${toolInfo?.name || "Unknown"}
- 描述：${toolInfo?.description || "No description"}
- 参数结构：${JSON.stringify(toolInfo?.inputSchema || {}, null, 2)}

**请提供：**
1. 错误原因分析
2. 针对该工具的具体使用指导
3. 改进后的任务指令（保持原始目标不变，但优化工具使用方式）
4. 预防类似错误的建议

请以JSON格式返回，包含以下字段：
{
  "analysis": "错误原因分析",
  "toolGuidance": "工具使用指导",
  "improvedInstructions": "改进后的指令",
  "preventionTips": "预防建议"
}
`;
  }

  /**
   * 生成纠正后的指令
   */
  private async generateCorrectedInstructions(
    correctionPrompt: string
  ): Promise<any> {
    if (!this.llm) return null;

    try {
      // 创建一个临时的LLM实例用于分析
      const analysisLLM = new ChatOpenAI(
        this.model,
        "你是一个专业的提示词优化专家，专门分析和改进AI工具使用指令。"
      );
      const response = await analysisLLM.chat(correctionPrompt);

      // 尝试解析JSON响应
      const jsonMatch = response.content.match(/\{[\s\S]*\}/);
      if (jsonMatch) {
        return JSON.parse(jsonMatch[0]);
      }

      return {
        analysis: "无法解析分析结果",
        toolGuidance: "请检查工具参数格式",
        improvedInstructions: this.originalPrompt,
        preventionTips: "确保JSON格式正确",
      };
    } catch (error) {
      process.stdout.write(`Prompt correction analysis failed: ${error}`);
      return null;
    }
  }

  /**
   * 应用纠正后的指令
   */
  private async applyCorrectedInstructions(corrections: any) {
    if (!corrections || !this.llm) return;

    process.stdout.write(`\n=== PROMPT CORRECTION RESULTS ===`);
    process.stdout.write(`错误分析: ${corrections.analysis}`);
    process.stdout.write(`工具指导: ${corrections.toolGuidance}`);
    process.stdout.write(`预防建议: ${corrections.preventionTips}`);
    process.stdout.write(`=================================\n`);

    // 向当前对话中注入改进的指令
    const improvementMessage = `
基于之前的错误分析，请注意以下改进指导：

**工具使用指导：**
${corrections.toolGuidance}

**改进后的任务指令：**
${corrections.improvedInstructions}

**预防建议：**
${corrections.preventionTips}

请根据以上指导重新执行任务，特别注意工具参数的正确格式。
`;

    // 将改进指令添加到对话历史中
    this.llm.appendToolResult("system-correction", improvementMessage);
  }

  /**
   * 启用或禁用提示词纠正功能
   */
  public setPromptCorrectionEnabled(enabled: boolean) {
    this.promptCorrectionEnabled = enabled;
  }

  /**
   * 设置最大失败次数阈值
   */
  public setMaxFailures(maxFailures: number) {
    this.maxFailures = maxFailures;
  }

  /**
   * 获取当前失败统计
   */
  public getFailureStats(): Map<string, number> {
    return new Map(this.failureCount);
  }

  /**
   * 重置失败计数
   */
  public resetFailureCount() {
    this.failureCount.clear();
  }
}
